/*
 *  $Id$
 *  Copyright [2010] [Panxiaobo] 
 *  Licensed under the Apache License, Version 2.0 (the "License"); 
 *  you may not use this file except in compliance with the License. 
 *  You may obtain a copy of the License at
 *  
 *     http://www.apache.org/licenses/LICENSE-2.0
 *     
 *  Unless required by applicable law or agreed to in writing, software 
 *  distributed under the License is distributed on an "AS IS" BASIS, 
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
 *  See the License for the specific language governing permissions and 
 *  limitations under the License. 
 */
package pxb.openid.rp;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.UnsupportedEncodingException;
import java.math.BigInteger;
import java.net.URLDecoder;
import java.security.KeyPair;
import java.util.Calendar;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.StringTokenizer;
import java.util.TreeMap;

import javax.crypto.interfaces.DHPrivateKey;
import javax.crypto.interfaces.DHPublicKey;
import javax.crypto.spec.DHParameterSpec;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import pxb.openid.Association;
import pxb.openid.AssociationImpl;
import pxb.openid.Constants;
import pxb.openid.OpenIdException;
import pxb.openid.util.Base64;
import pxb.openid.util.Dh;
import pxb.openid.util.Util;
import pxb.util.http.HttpConnector;
import pxb.util.http.HttpConnector.HttpResponse;
import pxb.util.http.impl.SimpleHttpConnectorImpl;

/**
 * @author Panxiaobo
 * 
 */
public class Associater implements Constants {
    private static final Logger log = LoggerFactory.getLogger(Associater.class);

    private static boolean isGoodValue(String value) {
        if ("&".equals(value) || ";".equals(value)) {
            return false;
        }
        // more tests here perchance
        return true;
    }

    private HttpConnector httpConnector;

    /**
	 * 
	 */
    public Associater() {
        this(new SimpleHttpConnectorImpl());
    }

    /**
     * @param httpConnector
     */
    public Associater(HttpConnector httpConnector) {
        super();
        this.httpConnector = httpConnector;
    }

    /**
     * 8.1. Association Session Request
     * 
     * @param services
     * @throws IOException
     * @throws ClientProtocolException
     */
    public Association associate(String server) {
        log.debug("start associate with server :{}", server);

        Date time = new Date();
        boolean notEncryption = false;
        boolean dhsha256 = true;
        // server.startsWith("https://");

        log.debug("notEncryption :{}", notEncryption);

        Map<String, String> aReq = new HashMap<String, String>();

        DHPrivateKey dhPrivateKey = null;

        // 8.1.1. Common Request Parameters
        aReq.put(OPENID_NS, OPENID_20_NAMESPACE);
        aReq.put(OPENID_MODE, "associate");
        aReq.put(OPENID_ASSOCIATION_TYPE, dhsha256 ? HMAC_SHA256 : HMAC_SHA1);
        if (notEncryption) {
            aReq.put(OPENID_SESSION_TYPE, NO_ENCRYPTION);
        } else {
            KeyPair aliceKpair = Dh.generateKeyPair(512);
            DHPublicKey alicePublicKey = (DHPublicKey) aliceKpair.getPublic();
            dhPrivateKey = (DHPrivateKey) aliceKpair.getPrivate();
            DHParameterSpec aliceParamSpec = alicePublicKey.getParams();
            BigInteger modules = aliceParamSpec.getP();
            BigInteger generator = aliceParamSpec.getG();
            BigInteger pubKey = alicePublicKey.getY();

            aReq.put(OPENID_SESSION_TYPE, dhsha256 ? DH_SHA256 : DH_SHA1);
            // 8.1.2. Diffie-Hellman Request Parameters
            aReq.put(OPENID_DH_CONSUMER_PUBLIC, Util.BigIntegerToBase64(pubKey));
            aReq.put(OPENID_DH_GENERATOR, Util.BigIntegerToBase64(generator));
            aReq.put(OPENID_DH_MODULUS, Util.BigIntegerToBase64(modules));
        }

        if (log.isDebugEnabled()) {
            log.debug("RP->OP associate request params:");
            log.debug("==================== Start");
            Map<String, String> map = new TreeMap<String, String>(aReq);
            for (Map.Entry<String, String> entry : map.entrySet()) {
                log.debug("{} = {}", entry.getKey(), entry.getValue());
            }
            log.debug("==================== End");
        }

        Map<String, String> aResp;
        try {
            aResp = doAssociate(server, aReq);
        } catch (IOException e) {
            throw new OpenIdException("Error Associater OpenID Provider", e);
        }
        if (log.isDebugEnabled()) {
            log.debug("OP->RP associate response params:");
            log.debug("==================== Start");
            Map<String, String> map = new TreeMap<String, String>(aResp);
            for (Map.Entry<String, String> entry : map.entrySet()) {
                log.debug("{} = {}", entry.getKey(), entry.getValue());
            }
            log.debug("==================== End");
        }

        if (aResp.containsKey("error")) {
            String error = aResp.get("error");
            String errorCode = aResp.get("error_code");
            throw new OpenIdException(String.format("OP Resturn ERROR:%s CODE:%s", error, errorCode));
        }

        byte[] macKey;
        if (notEncryption) {
            macKey = Base64.decodeBase64(aResp.get("mac_key").getBytes());
        } else {
            byte[] encodecMacKey = Base64.decodeBase64(aResp.get("enc_mac_key").getBytes());
            macKey = Util.xor(Util.sha(dhsha256 ? 256 : 1,
                    Dh.shared(dhPrivateKey, Util.Base64ToBigInteger(aResp.get("dh_server_public"))).toByteArray()),
                    encodecMacKey);
        }

        AssociationImpl association = new AssociationImpl();
        // type: String
        association.setId(aResp.get("assoc_handle"));
        association.setType(aResp.get("assoc_type"));
        association.setMac(macKey);

        Calendar calendar = Calendar.getInstance();
        calendar.setTime(time);
        calendar.add(Calendar.SECOND, Integer.parseInt(aResp.get("expires_in")));
        association.setExpireTime(calendar.getTime());
        log.debug("will expires at {}", calendar.getTime());
        return association;
    }

    protected Map<String, String> doAssociate(String server, Map<String, String> aReq) throws IOException {
        HttpResponse p = null;
        try {
            p = httpConnector.post(server, null, aReq);
            return read(p.getInputStream());
        } finally {
            if (p != null)
                p.close();
        }
    }

    public boolean isExpired(Association associate) {
        Date time = associate.getExpireTime();
        Date now = new Date();
        return time.before(now);
    }

    private Map<String, String> read(InputStream in) throws IOException {
        BufferedReader reader = new BufferedReader(new InputStreamReader(in));
        // line 1
        String l1 = reader.readLine();
        // line 2..n
        String ln = reader.readLine();
        if (ln == null) {// it is a query string
            return urlEncodedToMap(l1);
        }
        Map<String, String> map = new HashMap<String, String>();
        // put line 1
        int index = l1.indexOf(':');
        if (index != -1) {
            String name = l1.substring(0, index);
            String value = l1.substring(index + 1);
            if (name != null && value != null) {
                map.put(name, value);
            }
        }
        while (ln != null) {
            index = ln.indexOf(':');
            if (index != -1) {
                String name = ln.substring(0, index);
                String value = ln.substring(index + 1);
                if (name != null && value != null) {
                    map.put(name, value);
                }
            }
            ln = reader.readLine();
        }
        reader.close();
        return map;
    }

    /**
     * @param httpConnector
     *            the httpConnector to set
     */
    public void setHttpConnector(HttpConnector httpConnector) {
        this.httpConnector = httpConnector;
    }

    private Map<String, String> urlEncodedToMap(String query) throws UnsupportedEncodingException {
        Map<String, String> map = new HashMap<String, String>();
        StringTokenizer st = new StringTokenizer(query, "?&=;", true);
        String previous = null;
        while (st.hasMoreTokens()) {
            String current = st.nextToken();
            if ("?".equals(current) || "&".equals(current) || ";".equals(current)) {
                // ignore
            } else if ("=".equals(current)) {
                String name = URLDecoder.decode(previous, "UTF-8");
                if (st.hasMoreTokens()) {
                    String value = URLDecoder.decode(st.nextToken(), "UTF-8");
                    if (isGoodValue(value)) {
                        map.put(name, value);
                    }
                }
            } else {
                previous = current;
            }
        }
        return map;
    }
}
